#!/usr/bin/env python
# -*- coding: utf-8 -*-
__author__ = 'chenwei'
__version__ = 'v1.0.0'

from Crypto.Cipher import AES
from binascii import b2a_hex, a2b_hex
import pickle

def open_file(fn, md='rb'):  # 打开文件函数
    try:
        return open(fn, md)
    except IOError as e:
        raise e


def load_pickle(fn):#读取序列化文件。
    with open_file(fn) as f:
        try:
            enc_pkl_item = f.read()
            print(enc_pkl_item)
            passwd_key = "password_manager"
            pc = prpcrypt(passwd_key)  # 初始化密钥
            pkl_item = pc.decrypt_hex(enc_pkl_item)
            print(pkl_item)
            pkl_item = a2b_hex(pkl_item)
            return pickle.loads(pkl_item)
        except EOFError as e:
            pass
        except Exception as e:
            raise e


def dump_pickle(item,fn):#保存序列化文件。
    with open_file(fn, md='w+b') as f:
        try:
            pkl_item = pickle.dumps(item)
            passwd_key = "password_manager"
            pc = prpcrypt(passwd_key)  # 初始化密钥
            enc_pkl_item = pc.encrypt_hex(pkl_item)
            f.write(enc_pkl_item)
        except Exception as e:
            raise e
    return True





class prpcrypt():
    def __init__(self, key):
        self.key = key
        self.mode = AES.MODE_CBC

    # 加密函数，如果text不是16的倍数【加密文本text必须为16的倍数！】，那就补足为16的倍数
    def encrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
        length = 16
        count = len(text)
        add = length - (count % length)
        text = text + ('\0' * add)
        self.ciphertext = cryptor.encrypt(text)
        # 因为AES加密时候得到的字符串不一定是ascii字符集的，输出到终端或者保存时候可能存在问题
        # 所以这里统一把加密后的字符串转化为16进制字符串
        return str(b2a_hex(self.ciphertext),encoding="utf-8")

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        plain_text = cryptor.decrypt(a2b_hex(text))
        return str(plain_text,encoding = "utf-8").rstrip('\0')

    def encrypt_hex(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        # 这里密钥key 长度必须为16（AES-128）、24（AES-192）、或32（AES-256）Bytes 长度.目前AES-128足够用
        text2 = str(b2a_hex(text), encoding="utf-8")
        length = 16
        count = len(text2)
        add = length - (count % length)
        text2 = text2 + ('\0' * add)
        self.ciphertext = cryptor.encrypt(text2)
        return self.ciphertext

    # 解密后，去掉补足的空格用strip() 去掉
    def decrypt_hex(self, text):
        cryptor = AES.new(self.key, self.mode, self.key)
        plain_text = cryptor.decrypt(text)
        return str(plain_text,encoding = "utf-8").rstrip('\0')

if __name__ == '__main__':
    a = {'l1':{'passwd': '1111','delete_flag':1}}
    dump_pickle(a,'test.pkl')
    b = load_pickle('test.pkl')
    print(a)
    print(b)